package com.ds.lens.data.common.query;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.ElasticSearchDruidDataSourceFactory;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.ds.lens.data.common.query.es.MQL2ElasticSQLAdaptor;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.util.*;

public class ESQueryAdaptor implements QueryAdaptor<List<Map<String, Object>>> {
    Properties properties;
    DruidDataSource dds;

    public ESQueryAdaptor(Properties properties) {
        this.properties = properties;
    }

    @Override
    public void connect() throws Exception {
        dds = (DruidDataSource) ElasticSearchDruidDataSourceFactory.createDataSource(properties);
    }

    @Override
    public List<Map<String, Object>> query(String sql, MQLVariable... variables) throws Exception {
        sql = translate(sql, variables);
        try (Connection connection = dds.getConnection();
             PreparedStatement ps = connection.prepareStatement(sql);
             ResultSet resultSet = ps.executeQuery()) {
            ResultSetMetaData metaData = resultSet.getMetaData();
            List<Map<String, Object>> results = new ArrayList<>();
            while (resultSet.next()) {
//            String scrollId = resultSet.getString("_scroll_id");
                Map<String, Object> rowMap = new HashMap<>();
                for (int i = 1; i <= metaData.getColumnCount(); i++) {
                    String key = metaData.getColumnName(i);
                    Object value = resultSet.getObject(key);
                    rowMap.put(key, value);
                }
                results.add(rowMap);
            }
            return results;
        }
    }

    @Override
    public void close() throws Exception {
        dds.close();
    }

    private String translate(String sql, MQLVariable... variables) {
        List<SQLStatement> statements = SQLUtils.parseStatements(sql, null);
        SQLStatement statement = statements.get(0);
        SQLSelectStatement sqlSelectStatement = (SQLSelectStatement) statement;
        StringBuilder builder = new StringBuilder();
        MQL2ElasticSQLAdaptor translator = new MQL2ElasticSQLAdaptor(builder);
        sqlSelectStatement.accept(translator);
        sql = builder.toString();
        if (variables != null && variables.length > 0) {
            for (MQLVariable variable : variables) {
                sql = sql.replace(variable.getText(), variable.getValue());
            }
        }
        return sql;
    }

}
